import os  # 用于操作系统相关功能，如文件路径

import matplotlib.pyplot as plt
import numpy as np  # 导入 NumPy 库用于数组操作
import torch  # 导入 PyTorch 库，用于构建和训练神经网络
import torch.nn as nn  # 导入 PyTorch 的神经网络模块
import torch.optim as optim  # 导入优化器模块
from skimage import io  # 导入 skimage 用于图像读取
from torch.utils.data import Dataset, DataLoader  # 导入数据集和数据加载器
from torchvision import transforms, models  # 导入图像变换和预定义模型


# 1. 数据集定义
class CustomDataset(Dataset):
    def __init__(self, input_dir, mask_dir, transform=None):
        """
        初始化自定义数据集
        :param input_dir: 输入图像的目录
        :param mask_dir: 掩码图像的目录
        :param transform: 应用的图像变换
        """
        self.input_dir = input_dir
        self.mask_dir = mask_dir
        self.input_name = os.listdir(input_dir)  # 获取所有输入图像的文件名
        self.transform = transform

    def __len__(self):
        return len(self.input_name)  # 返回数据集的大小

    def __getitem__(self, idx):
        """
        获取指定索引的数据
        :param idx: 数据索引
        :return: 图像张量和对应的掩码张量
        """
        img_name = self.input_name[idx]
        img_path = os.path.join(self.input_dir, img_name)  # 构建图像路径
        mask_path = os.path.join(self.mask_dir, img_name)  # 构建掩码路径

        try:
            # 检查文件是否存在
            if not os.path.exists(img_path):
                raise FileNotFoundError(f"图像文件不存在: {img_path}")
            if not os.path.exists(mask_path):
                raise FileNotFoundError(f"掩码文件不存在: {mask_path}")

            # 读取图像和掩码
            image = io.imread(img_path)  # 读取图像
            mask = io.imread(mask_path, as_gray=True)  # 读取掩码，灰度图
            mask = mask.squeeze(0)  # 先去掉大小为1的维度
            mask = mask.reshape(image.shape[0], image.shape[1], 1)  # 在最后一个维度添加一个大小为1的维度

            # 检查图像和掩码是否成功读取
            if image is None:
                raise ValueError(f"无法读取图像文件: {img_path}")
            if mask is None:
                raise ValueError(f"无法读取掩码文件: {mask_path}")

            # 定义分割任务的图像变换
            if self.transform:
                image = self.transform(image)  # 应用图像变换

                # 对掩码应用不同的变换
                mask_transform = transforms.Compose([
                    transforms.ToTensor(),  # 将掩码转换为张量
                    transforms.Resize((256, 256)),  # 调整掩码大小为256x256
                ])
                mask = mask_transform(mask)  # 应用掩码变换
            # a,b = deepcopy(image), deepcopy(mask)
            # display_image_and_mask(a, b)
            return image.float(), mask.float()  # 返回图像张量和掩码张量

        except Exception as e:
            print(f"Error processing file {img_name}: {e}")
            # 返回一个全零的图像和掩码，或者根据需要处理
            return torch.zeros(3, 256, 256), torch.zeros(1, 256, 256)


# 2. 分割数据加载器定义（如果需要不同的变换）
# 这里我们已经在 CustomDataset 中处理了掩码的变换，因此不需要额外的分割数据加载器

# 3. 模型定义
class ResNetSegmentation(nn.Module):
    def __init__(self, pretrained=True, num_classes=1):
        """
        使用 ResNet-18 作为骨干网络的分割模型
        :param pretrained: 是否使用预训练的 ResNet-18 权重
        :param num_classes: 输出的类别数（单通道分割为1，多通道根据需要调整）
        """
        super(ResNetSegmentation, self).__init__()
        self.num_classes = num_classes

        # 加载预训练的 ResNet-18 模型
        self.backbone = models.resnet18()

        # 去掉 ResNet-18 的最后一个全连接层和平均池化层
        self.backbone = nn.Sequential(*list(self.backbone.children())[:-2])

        # 上采样模块
        self.upsample = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2),  # 从512通道上采样到256通道
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2),  # 从256通道上采样到128通道
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2),  # 从128通道上采样到64通道
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2),  # 从64通道上采样到32通道
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),

            nn.ConvTranspose2d(32, num_classes, kernel_size=2, stride=2),  # 从64通道上采样到输出通道数
        )

        # 激活函数
        self.sigmoid = nn.Sigmoid()  # 用于二值分割

    def forward(self, x):
        """
        前向传播
        :param x: 输入图像张量
        :return: 分割掩码张量
        """
        features = self.backbone(x)  # 提取特征
        x = self.upsample(features)  # 上采样恢复空间尺寸
        x = self.sigmoid(x)  # 应用 Sigmoid 激活函数，输出概率值
        return x


# 4. 训练与验证函数
def train_model(model, dataloader, criterion, optimizer, device, num_epochs=25):
    """
    训练模型
    :param model: 待训练的模型
    :param dataloader: 数据加载器
    :param criterion: 损失函数
    :param optimizer: 优化器
    :param device: 设备（CPU或GPU）
    :param num_epochs: 训练轮数
    :return: 训练好的模型
    """
    model.train()  # 设置模型为训练模式

    for epoch in range(num_epochs):
        print(f'第{epoch + 1}轮训练开始:')
        running_loss = 0.0  # 累计损失

        for batch_idx, (images, masks) in enumerate(dataloader):
            images = images.to(device)  # 将图像移至设备
            masks = masks.to(device)  # 将掩码移至设备

            optimizer.zero_grad()  # 清零梯度

            outputs = model(images)  # 前向传播
            loss = criterion(outputs, masks)  # 计算损失

            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数

            running_loss += loss.item() * images.size(0)  # 累加损失

            # 在每个batch后打印当前损失
            print(
                f'Epoch [{epoch + 1}/{num_epochs}], Batch [{batch_idx + 1}/{len(dataloader)}], Loss: {loss.item():.4f}')

        epoch_loss = running_loss / len(dataloader.dataset)  # 计算平均损失
        print(f'第{epoch + 1}轮训练结束, 平均Loss: {epoch_loss:.4f}')

        # 每个epoch结束后保存一次模型
        if (epoch + 1) % 5 == 0:
            torch.save(model.state_dict(), f'../pt_file/resnet_segmentation_epoch_{epoch + 1}.pt')
            print(f'模型已保存: resnet_segmentation_epoch_{epoch + 1}.pt')

    print('训练完成！')
    return model


def display_image_and_mask(image, mask):
    """
    显示输入图像和掩码，图像并排展示，操作前对输入数据进行拷贝。

    参数:
    - image: Tensor 或 NumPy 数组，输入图像，形状为 (3, H, W)
    - mask: Tensor 或 NumPy 数组，掩码图像，形状为 (1, H, W) 或 (H, W)
    """
    # 检查输入是否为Tensor，如果是，则转换为NumPy数组
    if torch.is_tensor(image):
        image = image.cpu().numpy().copy()  # 转换并拷贝
    else:
        image = np.copy(image)  # 拷贝NumPy数组

    if torch.is_tensor(mask):
        mask = mask.cpu().numpy().copy()  # 转换并拷贝
    else:
        mask = np.copy(mask)  # 拷贝NumPy数组

    # 转置图像数据以符合 (H, W, C) 的格式
    if image.shape[0] == 3:  # RGB图像
        image = image.transpose(1, 2, 0)  # 从 (3, H, W) 转换为 (H, W, 3)

    # 确保掩码是二维的
    if mask.ndim == 3 and mask.shape[0] == 1:  # 如果掩码是 (1, H, W)
        mask = mask.squeeze(0)  # 去除单一的颜色通道维度

    plt.figure(figsize=(12, 6))

    # 显示输入图像
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title("Input Image")
    plt.axis('off')

    # 显示掩码图像
    plt.subplot(1, 2, 2)
    plt.imshow(mask, cmap='gray')
    plt.title("Mask Image")
    plt.axis('off')

    plt.show()

    # 示例使用：
    # 假设image_tensor和mask_tensor是已加载的Tensor
    # display_image_and_mask(image_tensor, mask_tensor)


# 5. 主训练脚本
def main():
    # 数据集路径
    input_dir = '../unet_img/imgs/val'  # 输入图像目录
    mask_dir = '../unet_img/masks/val'  # 掩码图像目录

    # 定义图像变换（数据增强和预处理）
    transform = transforms.Compose([
        transforms.ToTensor(),  # 将图像转换为张量
        transforms.Resize((256, 256)),  # 调整图像大小为256x256
        transforms.Normalize(mean=[0.485, 0.456, 0.406],  # 使用ImageNet的均值
                             std=[0.229, 0.224, 0.225]),  # 使用ImageNet的标准差
    ])

    # 创建数据集和数据加载器
    dataset = CustomDataset(input_dir, mask_dir, transform=transform)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4)

    # 初始化设备（GPU 或 CPU）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f'使用设备: {device}')

    # 初始化模型
    model = ResNetSegmentation(pretrained=True, num_classes=1)  # 二值分割
    model = model.to(device)  # 将模型移至设备

    # 定义损失函数和优化器
    criterion = nn.BCELoss()  # 二值交叉熵损失函数
    optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam 优化器，学习率为0.001

    # 开始训练
    num_epochs = 25  # 设定训练的轮数
    trained_model = train_model(model, dataloader, criterion, optimizer, device, num_epochs=num_epochs)

    # 保存最终训练好的模型
    torch.save(trained_model.state_dict(), '../pt_file/resnet_segmentation_final.pt')
    print('训练完成，模型已保存为 resnet_segmentation_final.pt')


if __name__ == "__main__":
    main()
